# 导入tushare
import tushare as ts
# 初始化pro接口
pro = ts.pro_api('f31f00fef45ea66da46b2f020f5599fe139920e0b80afa76debe7b47')

# 拉取数据
df = pro.daily(**{
    "ts_code": "",
    "trade_date": "",
    "start_date": 20250320,
    "end_date": 20250330,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
# print(df)

import pandas as pd

# 假设df是已经获取到的数据
# 计算MA5和MA10
df['MA5'] = df.groupby('ts_code')['close'].rolling(window=5).mean().reset_index(0,drop=True)
df['MA10'] = df.groupby('ts_code')['close'].rolling(window=10).mean().reset_index(0,drop=True)

# 计算RSI
def calculate_rsi(series, period=14):
    delta = series.diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

df['RSI'] = df.groupby('ts_code')['close'].apply(calculate_rsi)

# 按股票代码和交易日期排序
df = df.sort_values(by=['ts_code', 'trade_date'])

# 标记5日线上穿10日线的点
df['MA5_cross_MA10'] = (df['MA5'] > df['MA10']) & (df.groupby('ts_code')['MA5'].shift(1) < df.groupby('ts_code')['MA10'].shift(1))

# 提取买点股票数据
buy_points = df[df['MA5_cross_MA10']]
print(df.head(5))